# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: xguo
import time
import json
import google.protobuf.json_format as json_format

import coin.proto.coin_order_enums_pb2 as coin_order
from coin.exchange.base.kr_rest.private_client_base import PrivateClientBase
from coin.exchange.okex_futures.kr_rest.futures_product import OkexFuturesProduct
from coin.exchange.okex.kr_rest.currency import OkexCurrency
from coin.proto.coin_query_pb2 import (
    AccountBalance,
    AccountPosition,
    CurrencyBalance,
    ProductOrderBundle,
    ProductOrderElement,
    ProductPosition,
    ProductFillElement,
    ProductFillBundle,
)
from coin.exchange.okex_futures_v3.kr_rest.native_private_client import (
    OkexFuturesNativePrivateClient, OkexFuturesOrderStatus)
from coin.exchange.okex_futures_v3.kr_rest.product_info import okex_futures_product_info    
from cc.appcoin2.support.order.py_parser import PyOrderParser
from coin2.exchange.order.system_pb2 import OrderSystemConfig
from coin2.exchange.symbology_pb2 import ProductConfig

def convert_okex_order_type_to_proto_side(okex_order_type):
  okex_order_type = int(okex_order_type)
  if okex_order_type == 1:
    side = coin_order.BUY_OPEN_ORDER
  elif okex_order_type == 2:
    side = coin_order.SELL_OPEN_ORDER
  elif okex_order_type == 3:
    side = coin_order.BUY_CLOSE_ORDER
  elif okex_order_type == 4:
    side = coin_order.SELL_CLOSE_ORDER
  else:
    raise ValueError('Invalid okex order type %s' % okex_order_type)
  return side


def convert_okex_exec_type_to_proto_fill_type(okex_exec_type):
  if okex_exec_type == 'T':
    fill_type = coin_order.TAKER_FILL_TYPE
  elif okex_exec_type == 'M':
    fill_type = coin_order.MAKER_FILL_TYPE
  else:
    raise ValueError('Invalid okex exec type %s' % okex_exec_type)
  return fill_type


class OkexFuturesPrivateParser(object):
  @staticmethod
  def parse_cross_margin_balance(msg):
    """"
    {
    "result":true,  # this field is removed
    "info":{
        "btc":{
            "equity":"0.0005297",
            "margin":"0",
            "margin_mode":"crossed",
            "margin_ratio":"10000",
            "realized_pnl":"0",
            "total_avail_balance":"0.0005297",
            "unrealized_pnl":"0"
        },
        "btg":{
            "equity":"0",
            "margin":"0",
            "margin_mode":"crossed",
            "margin_ratio":"10000",
            "realized_pnl":"0",
            "total_avail_balance":"0",
            "unrealized_pnl":"0"
        },
        "etc":{
            "equity":"0",
            "margin":"0",
            "margin_mode":"crossed",
            "margin_ratio":"10000",
            "realized_pnl":"0",
            "total_avail_balance":"0",
            "unrealized_pnl":"0"
        },
        "bch":{
            "equity":"0",
            "margin":"0",
            "margin_mode":"crossed",
            "margin_ratio":"10000",
            "realized_pnl":"0",
            "total_avail_balance":"0",
            "unrealized_pnl":"0"
        },
     }
    }
    """
    balance_list = []
    for native_currency, balance_info in msg['info'].items():
      assert balance_info['margin_mode'] == 'crossed', balance_info
      currency = OkexCurrency.FromStrNativeCurrency(native_currency)
      total = float(balance_info['equity'])
      hold = float(balance_info['margin'])
      available = total - hold
      balance = CurrencyBalance(
          currency=currency.currency,
          total=total,
          available=available,
          hold=hold,
      )
      balance_list.append(balance)
    return AccountBalance(
        exchange='Okex',
        each_balance=balance_list,
    )

  @staticmethod
  def parse_margin_balance(msg):
    balance_list = []
    currency_list = []
    for _, balance_info in msg['info'].items():
      assert balance_info['margin_mode'] == 'crossed', balance_info
      currency = OkexCurrency.FromStrNativeCurrency(balance_info['currency'])
      total = float(balance_info['equity'])
      hold = float(balance_info['margin'])
      available = total - hold
      if currency in currency_list:
        index_num = currency_list.index(currency)
        balance_list[index_num].total += total
        balance_list[index_num].available += available
        balance_list[index_num].hold += hold
      else:
        balance = CurrencyBalance(
            currency=currency.currency,
            total=total,
            available=available,
            hold=hold,
        )
        currency_list.append(currency)
        balance_list.append(balance)
    return AccountBalance(
        exchange='Okex',
        each_balance=balance_list,
    )

  @staticmethod
  def parse_position(msg):
    """
    {
      "result":true,
      "holding":[
          {
              "long_qty":"1",
              "long_avail_qty":"1",
              "long_avg_cost":"5.404",
              "long_settlement_price":"5.404",
              "realised_pnl":"-0.584",
              "short_qty":"0",
              "short_avail_qty":"0",
              "short_avg_cost":"5.331",
              "short_settlement_price":"5.331",
              "liquidation_price":"0.796",
              "instrument_id":"EOS-USD-181026",
              "leverage":"10",
              "created_at":"2018-10-19T10:39:59.0Z",
              "updated_at":"2018-10-24T03:46:25.0Z"
          }
      ],
      "margin_mode":"crossed"
    }
    """
    if 'result' in msg:
      assert msg['result'], msg

    if 'holding' in msg:
      if msg['holding']:
        assert len(msg['holding']) == 1, msg
        pos_info_list = msg['holding'][0]
      else:
        pos_info_list = []
    else:
      pos_info_list = msg['data']

    position_list = []
    for pos_info in pos_info_list:
      product = OkexFuturesProduct.FromStrNativeProduct(pos_info['instrument_id'])
      long_position = float(pos_info['long_qty'])
      short_position = -float(pos_info['short_qty'])
      long_avg_cost = float(pos_info['long_avg_cost'])
      short_avg_cost = float(pos_info['short_avg_cost'])
      available_long_position = float(pos_info['long_avail_qty'])
      available_short_position = -float(pos_info['short_avail_qty'])
      realized_pnl = float(pos_info['realised_pnl'])
      net_position = long_position + short_position
      position = ProductPosition(
          symbol=product.symbol,
          native_symbol=product.native_symbol,
          long_position=long_position,
          available_long_position=available_long_position,
          short_position=short_position,
          available_short_position=available_short_position,
          net_position=net_position,
          long_avg_entry_price=long_avg_cost,
          short_avg_entry_price=short_avg_cost,
          realized_pnl=realized_pnl,
      )
      position_list.append(position)

    return AccountPosition(exchange='Okex', market_type='Futures', each_position=position_list)

  @staticmethod
  def parse_ws_product_position(pos_msg):
    assert pos_msg['table'] == 'futures/position', pos_msg
    assert len(pos_msg['data']) == 1, pos_msg
    pos_info = pos_msg['data'][0]
    product = OkexFuturesProduct.FromStrNativeProduct(pos_info['instrument_id'])
    long_position = float(pos_info['long_qty'])
    short_position = -float(pos_info['short_qty'])
    long_avg_cost = float(pos_info['long_avg_cost'])
    short_avg_cost = float(pos_info['short_avg_cost'])
    available_long_position = float(pos_info['long_avail_qty'])
    available_short_position = -float(pos_info['short_avail_qty'])
    realized_pnl = float(pos_info['realised_pnl'])
    net_position = long_position + short_position

    position = ProductPosition(
        symbol=product.symbol,
        native_symbol=product.native_symbol,
        long_position=long_position,
        available_long_position=available_long_position,
        short_position=short_position,
        available_short_position=available_short_position,
        net_position=net_position,
        long_avg_entry_price=long_avg_cost,
        short_avg_entry_price=short_avg_cost,
        realized_pnl=realized_pnl,
    )
    return position

  @staticmethod
  def parse_orders(msg, product):
    """
    {
    "result":true,
    "order_info":[
        {
            "instrument_id":"EOS-USD-181026",
            "size":"1",
            "timestamp":"2018-10-23T07:51:06.0Z",
            "filled_qty":"0",
            "fee":"0",
            "order_id":"1674121456333824",
            "price":"1000000",
            "price_avg":"0",
            "status":"0",
            "type":"2",
            "contract_val":"10",
            "leverage":"10"
        },
        {
            "instrument_id":"EOS-USD-181026",
            "size":"1",
            "timestamp":"2018-10-23T07:58:41.0Z",
            "filled_qty":"0",
            "fee":"0",
            "order_id":"1674151290181632",
            "price":"1000000",
            "price_avg":"0",
            "status":"0",
            "type":"2",
            "contract_val":"10",
            "leverage":"10"
        }
    ]
    }
    """
    assert msg['result'], msg
    order_list = []
    for order_info in msg['order_info']:
      native_symbol = order_info['instrument_id']
      if native_symbol != product.native_symbol:
        return

      order = ProductOrderElement(
          symbol=product.symbol,
          order_id=order_info['order_id'],
          price=float(order_info['price']),
          qty=float(order_info['size']),
          side=convert_okex_order_type_to_proto_side(order_info['type']),
      )
      order_list.append(order)
    return ProductOrderBundle(
        exchange='Okex',
        symbol=product.symbol,
        each_order=order_list,
    )

  @staticmethod
  def parse_fills(order_list_msg, fill_msgs, product):
    """
    Order list is the same format as orders, but always status=2.
    fill update format:
    [
    {
      "trade_id":"2047624468856838",
      "instrument_id":"LTC-USD-181228",
      "price":"27.942",
      "order_qty":"1",
      "order_id":"2047624459797504",
      "created_at":"2018-12-28T06:57:49.0Z",
      "exec_type":"T",
      "fee":"-0.00010737",
      "side":"sell"
    }
    ]
    """
    assert order_list_msg['order_info'], order_list_msg
    fill_list = []
    for index, order_info in enumerate(order_list_msg['order_info']):
      native_symbol = order_info['instrument_id']
      if native_symbol != product.native_symbol:
        return

      assert int(order_info['status']) == OkexFuturesOrderStatus.FILLED
      fill_msg = fill_msgs[index]
      assert len(fill_msg) >= 1
      fill_msg = fill_msg[0]
      prod = OkexFuturesProduct.FromStrNativeProduct(fill_msg['instrument_id'])
      assert prod == product
      fee = -float(fill_msg['fee'])
      fill_type = convert_okex_exec_type_to_proto_fill_type(fill_msg['exec_type'])

      fill = ProductFillElement(symbol=product.symbol,
                                order_id=order_info['order_id'],
                                price=float(fill_msg['price']),
                                qty=float(fill_msg['order_qty']),
                                side=convert_okex_order_type_to_proto_side(order_info['type']),
                                fee=fee,
                                fee_currency=product.base.symbol,
                                fill_type=fill_type,
                                contract_val=float(order_info['contract_val']))
      fill_list.append(fill)
    return ProductFillBundle(
        exchange='Okex',
        symbol=product.symbol,
        each_fill=fill_list,
    )


class OkexFuturesPrivateClient(PrivateClientBase):
  ProductType = OkexFuturesProduct

  def __init__(self, key_file):
    super().__init__()
    self.nprivc = OkexFuturesNativePrivateClient(key_file=key_file)
    pi_bundle = okex_futures_product_info._product_infos
    assert pi_bundle.HasField('mea')
    oss_config = OrderSystemConfig(
        mea=pi_bundle.mea,
        products=ProductConfig(norms=[pi.symbol for pi in pi_bundle.product_infos]))
    self._parser = PyOrderParser.from_mea_str(pi_bundle.mea)
    self._parser.add_product(oss_config.SerializeToString())

  def query_account_balance_impl(self):
    update = self.nprivc.query_accounts()
    msg = self._parser.parse_account_balance(json.dumps(update.msg))
    account_balance = AccountBalance()
    json_format.Parse(msg, account_balance)
    update.msg = account_balance
    return update

  def query_account_position_impl(self):
    update = self.nprivc.query_positions()
    msg = self._parser.parse_account_position(json.dumps(update.msg))
    account_position = AccountPosition()
    json_format.Parse(msg, account_position)
    update.msg = account_position
    return update

  def query_list_orders_impl(self, product):
    update = self.nprivc.query_order_list(status=OkexFuturesOrderStatus.OPEN,
                                          instrument_id=product.native_symbol)
    update.msg = OkexFuturesPrivateParser.parse_orders(update.msg, product)
    return update

  def query_fills_impl(self, product):
    update = self.nprivc.query_order_list(status=OkexFuturesOrderStatus.FILLED,
                                          instrument_id=product.native_symbol,
                                          limit=50)
    fill_updates = []
    for order in update.msg['order_info']:
      native_symbol = order['instrument_id']
      if native_symbol != product.native_symbol:
        continue
      order_id = str(order['order_id'])
      fill_update = self.nprivc.query_transaction_details(order_id=order_id,
                                                          instrument_id=product.native_symbol)
      time.sleep(0.2)  # rate limit 10 times per 2 seconds
      fill_updates.append(fill_update.msg)
    update.msg = OkexFuturesPrivateParser.parse_fills(update.msg, fill_updates, product)
    return update

  def query_submit_impl(self, product, side, price, qty):
    raise NotImplementedError()

  def query_cancel_impl(self, product, order_id):
    raise NotImplementedError()

  def query_cancel_product_impl(self, product):
    raise NotImplementedError()

  def query_account_fills_impl(self):
    raise NotImplementedError()

  def query_account_list_orders_impl(self):
    raise NotImplementedError()

if __name__ == '__main__':
  pv = OkexFuturesPrivateClient('../../coin_key/view/prestoinvt07/OKEx/view_v3_key.json')
  # print(pv.query_account_balance_impl().msg)
  # print(pv.query_account_position_impl().msg)